# 此函数用于光纤信号的衰减
# 传入参数：
#   signal:信号
#   alpha:衰减
#
# 输出参数：
#   sig：衰减后的信号，时域
from numpy import exp
from numpy.fft import fft, fftshift, ifftshift, ifft
import numpy as np


def signal_attenuation(signal, Signal, Fiber_para, PMD_para):
    # [1/km]
    alpha_km = Fiber_para["alpha_km"]
    polar_n = Signal["polarization"]
    # [km]
    deltaz = Fiber_para["step_size"]
    PMD = PMD_para["PMD"]
    steptype = Fiber_para['ssfm_type']

    if polar_n == 2:
        data = signal
        data_x = data[0:, 0]
        data_y = data[0:, 1]
        # print('衰减前功率')
        # print(np.mean(np.abs(data_x) ** 2))
        data_fft_x = fftshift(fft(data_x))
        data_fft_y = fftshift(fft(data_y))
        # # 衰减
        data_fft_x = data_fft_x * exp((-alpha_km/2)*(deltaz/2))
        data_fft_y = data_fft_y * exp((-alpha_km/2)*(deltaz/2))
        # X 偏振
        data_x = ifft(ifftshift(data_fft_x))
        # Y 偏振
        data_y = ifft(ifftshift(data_fft_y))
        # 衰减
        # data_x = data_x * np.exp((-alpha_km / 2) * (deltaz / 2))
        # data_y = data_y * np.exp((-alpha_km / 2) * (deltaz / 2))

        # print('衰减后功率')
        # print(np.mean(np.abs(data_x) ** 2))

        data_x = data_x.reshape(-1, 1)
        data_y = data_y.reshape(-1, 1)
        sig = np.concatenate((data_x, data_y), axis=1)
    else:
        data = signal
        # data_fft = fftshift(fft(data))
        # # 衰减
        # data_fft = data_fft * exp(-alpha_km/2*deltaz)
        # data_t = ifft(ifftshift(data_fft))
        data = data * exp(-alpha_km / 2 * deltaz / 2)
        sig = data
    return sig
